package io.kotest.property.arbitrary

import io.kotest.property.Arb
import io.kotest.property.PropertyTesting
import io.kotest.property.RTree
import io.kotest.property.RandomSource
import io.kotest.property.Sample
import io.kotest.property.filter

/**
 * Returns a new [Arb] which takes its elements from the receiver and filters them using the supplied
 * predicate. This gen will continue to request elements from the underlying gen until one satisfies
 * the predicate.
 */
fun <A> Arb<A>.filter(predicate: (A) -> Boolean): Arb<A> = trampoline { sampleA ->
   object : Arb<A>() {
      override fun edgecase(rs: RandomSource): A? =
         sequenceOf(sampleA.value)
            .plus(generateSequence { this@filter.edgecase(rs) })
            .take(PropertyTesting.maxFilterAttempts)
            .filter(predicate)
            .firstOrNull()

      override fun sample(rs: RandomSource): Sample<A> {
         val sample = sequenceOf(sampleA).plus(this@filter.samples(rs)).filter { predicate(it.value) }.first()
         return Sample(sample.value, sample.shrinks.filter(predicate) ?: RTree({ sample.value }))
      }
   }
}

/**
 * @return a new [Arb] by filtering this arbs output by the negated function [f]
 */
fun <A> Arb<A>.filterNot(f: (A) -> Boolean): Arb<A> = filter { !f(it) }

/**
 * Create a new [Arb] by keeping only instances of B generated by this gen.
 * This is useful if you have a type hierarchy and only want to retain
 * a particular subtype.
 */
@Suppress("UNCHECKED_CAST")
inline fun <reified B> Arb<*>.filterIsInstance(): Arb<B> = filter { it is B } as Arb<B>
